package cz.semjob.learning.data;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;

import org.apache.log4j.Logger;

import cz.semjob.document.AnnotatedDocument;
import cz.semjob.document.annotation.metadata.Entity;
import cz.semjob.document.annotation.metadata.Relation;
import cz.semjob.learning.instances.LabeledRelationInstance;
import cz.semjob.learning.instances.LabeledTextInstance;
import cz.semjob.learning.instances.creator.ITextInstanceCreator;
import cz.semjob.learning.relation.IRelationInstanceCreator;

public class DataSetBatchCreator {

	private static Logger logger = Logger.getLogger(DataSetBatchCreator.class);

	private Map<Entity, ITextInstanceCreator> textInstanceCreators;

	private Map<String, IRelationInstanceCreator> relationInstanceCreators;

	public List<DataSet> createDataSets(List<AnnotatedDocument> documents,
			List<Entity> entities) throws TaxonomyLearningException {
		// File exportDirectory = new File(exportDirectoryPath);
		// if (!exportDirectory.isDirectory()) {
		// throw new
		// TaxonomyLearningException("Error while processing data set batch: '"
		// + exportDirectoryPath + "' is not a directory");
		// }
		List<DataSetCollector> dataSetCollectors = new ArrayList<DataSetCollector>();
		for (Entity entity : entities) {
			DataSetCollector entityCreator = new DataSetCollector(
					entity.getName(), textInstanceCreators.get(entity));
			dataSetCollectors.add(entityCreator);
		}

		for (AnnotatedDocument annotatedDocument : documents) {
			logger.info("Creating data set from document: "
					+ annotatedDocument.getDocumentInstance().getFileName());
			for (Entity entity : entities) {
				DataSetCollector entityDataSetCollector = dataSetCollectors
						.get(entities.indexOf(entity));
				if (annotatedDocument.containsAnnotationsForEntity(entity)) {
					fillDocumentInstancesForEntity(annotatedDocument, entity,
							entityDataSetCollector);

				}
			}
		}

		List<DataSet> dataSets = new ArrayList<DataSet>();
		for (DataSetCollector dataSetCollector : dataSetCollectors) {
			dataSets.add(dataSetCollector.getDataSet());
		}
		return dataSets;

	}

	public RelationDataSet createRelationDataSet(
			List<AnnotatedDocument> documents, String relationType,
			List<Relation> relations) throws TaxonomyLearningException {

		RelationDataSetCollector relationDataSetCollector = new RelationDataSetCollector(
				relationType, relationInstanceCreators.get(relationType),
				relations);

		for (AnnotatedDocument annotatedDocument : documents) {
			logger.info("Creating data set from document: "
					+ annotatedDocument.getDocumentInstance().getFileName());
			for (Relation relation : relations) {
				// RelationDataSetCollector relationDataSetCollector =
				// dataSetCollectors
				// .get(relations.indexOf(relation));
				if (annotatedDocument.containsAnnotationsForRelation(relation)) {
					fillDocumentInstancesForRelation(annotatedDocument,
							relation, relationDataSetCollector, relationType);

				}
			}
			// fillDocumentInstancesForRelation(annotatedDocument,
			// Relation.getNoneRelation(), relationDataSetCollector,
			// relationType);
		}

		RelationDataSet dataSet = relationDataSetCollector.getDataSet();
		return dataSet;

	}

	private void fillDocumentInstancesForEntity(
			AnnotatedDocument annotatedDocument, Entity entity,
			DataSetCollector entityDataSetCollector)
			throws TaxonomyLearningException {
		List<LabeledTextInstance> positiveInstances = textInstanceCreators.get(
				entity).createLabeledInstancesForEntity(annotatedDocument,
				entity, true);
		entityDataSetCollector.addPositiveInstances(positiveInstances);

		List<LabeledTextInstance> negativeInstances = textInstanceCreators.get(
				entity).createLabeledInstancesForEntity(annotatedDocument,
				entity, false);
		entityDataSetCollector.addNegativeInstances(negativeInstances);
	}

	private void fillDocumentInstancesForRelation(
			AnnotatedDocument annotatedDocument, Relation relation,
			RelationDataSetCollector relationDataSetCollector,
			String relationType) throws TaxonomyLearningException {
		List<LabeledRelationInstance> positiveInstances = relationInstanceCreators
				.get(relationType).createRelationLabeledInstances(
						annotatedDocument, relation);
		if (positiveInstances != null) {
			relationDataSetCollector.addInstances(positiveInstances, relation);
		}
	}

	public DataSet createDataSet(List<AnnotatedDocument> documents,
			Entity entity) throws TaxonomyLearningException {
		DataSetCollector entityDataSetCollector = new DataSetCollector(
				entity.getName(), textInstanceCreators.get(entity));
		for (AnnotatedDocument annotatedDocument : documents) {
			fillDocumentInstancesForEntity(annotatedDocument, entity,
					entityDataSetCollector);
		}
		return entityDataSetCollector.getDataSet();
	}

	public Map<Entity, ITextInstanceCreator> getTextInstanceCreators() {
		return textInstanceCreators;
	}

	public void setTextInstanceCreators(
			Map<Entity, ITextInstanceCreator> textInstanceCreators) {
		this.textInstanceCreators = textInstanceCreators;
	}

	public void setRelationInstanceCreators(
			Map<String, IRelationInstanceCreator> relationInstanceCreators) {
		this.relationInstanceCreators = relationInstanceCreators;
	}

}
